clc;
clear;

tic;
startTime = cputime;
pde = process_pde_t(1);
vector_h = [1/4 1/8 1/16];
iter_L = 5;

u_Linf = zeros(size(vector_h,2),1);
u_L2 = zeros(size(vector_h,2),1);
u_H1 = zeros(size(vector_h,2),1);
p_Linf = zeros(size(vector_h,2),1);
p_L2 = zeros(size(vector_h,2),1);
p_H1 = zeros(size(vector_h,2),1);

for i = 1:size(vector_h,2)
    h = vector_h(i);
    mesh = process_mesh(pde, h, "P2");

    O1 = sparse(mesh.Np, mesh.Np);
    O2 = sparse(mesh.Nb, mesh.Np);
    O3 = sparse(mesh.Nb, mesh.Nb);
    O4 = zeros(mesh.Np,1);
    A1 = assemble_stiffness_matrix(mesh, "P2", "dx", "P2", "dx", pde.nu);
    A2 = assemble_stiffness_matrix(mesh, "P2", "dy", "P2", "dy", pde.nu);
    A3 = assemble_stiffness_matrix(mesh, "P2", "dx", "P2", "dy", pde.nu);
    A5 = assemble_stiffness_matrix(mesh, "P1", "x", "P2", "dx", -1);
    A6 = assemble_stiffness_matrix(mesh, "P1", "x", "P2", "dy", -1);
    A = [2*A1+A2, A3, A5;
        A3', 2*A2+A1, A6;
        A5', A6', O1];
    Me = assemble_mass_matrix(mesh);
    M = [Me, O3, O2;
        O3, Me, O2;
        O2', O2', O1];
    U1 = [pde.u1(mesh.node(:,1),mesh.node(:,2),0);
        pde.u2(mesh.node(:,1),mesh.node(:,2),0);
        pde.p(mesh.node(1:mesh.Np,1),mesh.node(1:mesh.Np,2),0)];
    dt = 8*h^3;
    Nt = (pde.domain(6) - pde.domain(5))/dt;
    for m = 0:Nt-1
        t2 = (m+1) * dt;
        [b1, b2] = assemble_load_vector(mesh, pde, t2);
        bt2 = [b1;
            b2;
            O4];
        U_old = U1;
        for l = 1:iter_L
            AN1 = assemble_nonlinear_matrix(mesh, "u1", "dx", "P2", "x", "P2", "x", U_old);
            AN2 = assemble_nonlinear_matrix(mesh, "u1", "x", "P2", "dx", "P2", "x", U_old);
            AN3 = assemble_nonlinear_matrix(mesh, "u2", "x", "P2", "dy", "P2", "x", U_old);
            AN4 = assemble_nonlinear_matrix(mesh, "u1", "dy", "P2", "x", "P2", "x", U_old);
            AN5 = assemble_nonlinear_matrix(mesh, "u2", "dx", "P2", "x", "P2", "x", U_old);
            AN6 = assemble_nonlinear_matrix(mesh, "u2", "dy", "P2", "x", "P2", "x", U_old);
            AN = [AN1+AN2+AN3, AN4, O2;
                AN5, AN6+AN2+AN3, O2;
                O2', O2', O1];
            bN1 = assemble_nonlinear_vector(mesh, "u1", "x", "u1", "dx", "P2", "x", U_old);
            bN2 = assemble_nonlinear_vector(mesh, "u2", "x", "u1", "dy", "P2", "x", U_old);
            bN3 = assemble_nonlinear_vector(mesh, "u1", "x", "u2", "dx", "P2", "x", U_old);
            bN4 = assemble_nonlinear_vector(mesh, "u2", "x", "u2", "dy", "P2", "x", U_old);
            bN = [bN1+bN2;
                bN3+bN4;
                O4];
            Al = M./dt + A + AN;
            bl = (M./dt)*U_old + bt2 + bN;
            [Al, bl] = process_boundary_t(Al, bl, mesh, pde, t2);
            Al(2*mesh.Nb+1,:) = 0;
            Al(2*mesh.Nb+1,2*mesh.Nb+1) = 1;
            bl(2*mesh.Nb+1,1) = pde.p(pde.domain(1), pde.domain(3), t2);
            U2 = Al\bl;
            U_old = U2;
        end
        U1 = U2;
    end

    [u_Linf(i,1), p_Linf(i,1)] = process_error_t(mesh, pde, U2, "Linf");
    [u_L2(i,1), p_L2(i,1)] = process_error_t(mesh, pde, U2, "L2");
    [u_H1(i,1), p_H1(i,1)] = process_error_t(mesh, pde, U2, "H1");

    fprintf("已完成\t%d/%d\n", i, size(vector_h,2));
end

process_rate(u_Linf, vector_h);
process_rate(u_L2, vector_h);
process_rate(u_H1, vector_h);
process_rate(p_Linf, vector_h);
process_rate(p_L2, vector_h);
process_rate(p_H1, vector_h);

elapsedTOCTime = toc;
elapsedCPUTime = cputime - startTime;
disp(["TOC time(s)", num2str(elapsedTOCTime)]);
disp(["CPU time(s)", num2str(elapsedCPUTime)]);

%% assemble_stiffness_matrix
function matrix_A = assemble_stiffness_matrix(mesh, B1_type, D1_type, B2_type, D2_type, cor)
basis_map = containers.Map(...
    {'P1', 'P1b', 'P2'},...
    {struct('Nb', mesh.Np, 'Nlb',3),...
    struct('Nb', mesh.Nb, 'Nlb',4),...
    struct('Nb', mesh.Nb, 'Nlb',6)});
B1 = basis_map(B1_type);
Nb1 = B1.Nb();
Nlb1 = B1.Nlb();
B2 = basis_map(B2_type);
Nb2 = B2.Nb();
Nlb2 = B2.Nlb();
rows = zeros(2*Nlb1*Nlb2,1);
cols = rows;
values = rows;
flag = 1;
for n = 1:mesh.Ne
    E = mesh.node(mesh.elem(n,:),:);
    V = E(1:3,:);
    [gauss_points_loc, gauss_weights_loc] = gauss_lengdre_loc(V, "2d");
    phi1 = basis_function_loc(V, gauss_points_loc(:,1), gauss_points_loc(:,2), B1_type, D1_type);
    phi2 = basis_function_loc(V, gauss_points_loc(:,1), gauss_points_loc(:,2), B2_type, D2_type);
    for i = 1:Nlb2
        for j = 1:Nlb1
            f = cor * phi1(:,j) .* phi2(:,i);
            r = gauss_weights_loc * f;
            rows(flag) = mesh.elem(n,i);
            cols(flag) = mesh.elem(n,j);
            values(flag) = r;
            flag = flag + 1;
        end
    end
end
matrix_A = sparse(rows, cols, values, Nb2, Nb1);
end

%% assemble_mass_matrix
function matrix_M = assemble_mass_matrix(mesh)
Nb1 = mesh.Nb;
Nlb1 = mesh.Nlb;
Nb2 = Nb1;
Nlb2 = Nlb1;
rows = zeros(2*Nlb1*Nlb2,1);
cols = rows;
values = rows;
flag = 1;
for n = 1:mesh.Ne
    E = mesh.node(mesh.elem(n,:),:);
    V = E(1:3,:);
    [gauss_points_loc, gauss_weights_loc] = gauss_lengdre_loc(V, "2d");
    phi = basis_function_loc(V, gauss_points_loc(:,1), gauss_points_loc(:,2), mesh.basis_type, "x");
    for i = 1:Nlb2
        for j = 1:Nlb1
            f = phi(:,j) .* phi(:,i);
            r = gauss_weights_loc * f;
            rows(flag) = mesh.elem(n,i);
            cols(flag) = mesh.elem(n,j);
            values(flag) = r;
            flag = flag + 1;
        end
    end
end
matrix_M = sparse(rows, cols, values, Nb2, Nb1);
end

%% assemble_load_vector
function [vector_b1, vector_b2] = assemble_load_vector(mesh, pde, t)
vector_b1 = zeros(mesh.Nb,1);
vector_b2 = vector_b1;
for n = 1:mesh.Ne
    E = mesh.node(mesh.elem(n,:),:);
    V = E(1:3,:);
    [gauss_points_loc, gauss_weights_loc] = gauss_lengdre_loc(V, "2d");
    phi = basis_function_loc(V, gauss_points_loc(:,1), gauss_points_loc(:,2), mesh.basis_type, "x");
    ff1 = pde.f1(gauss_points_loc(:,1), gauss_points_loc(:,2), t);
    ff2 = pde.f2(gauss_points_loc(:,1), gauss_points_loc(:,2), t);
    for i = 1:mesh.Nlb
        f1 = ff1 .* phi(:,i);
        f2 = ff2 .* phi(:,i);
        r1 = gauss_weights_loc * f1;
        r2 = gauss_weights_loc * f2;
        vector_b1(mesh.elem(n,i),1) = vector_b1(mesh.elem(n,i),1) + r1;
        vector_b2(mesh.elem(n,i),1) = vector_b2(mesh.elem(n,i),1) + r2;
    end
end
end

%% assemble_nonlinear_matrix
function matrix_A = assemble_nonlinear_matrix(mesh, u_type, du_type, B1_type, D1_type, B2_type, D2_type, U_old)
basis_map = containers.Map(...
    {'P1', 'P1b', 'P2'},...
    {struct('Nb', mesh.Np, 'Nlb',3),...
    struct('Nb', mesh.Nb, 'Nlb',4),...
    struct('Nb', mesh.Nb, 'Nlb',6)});
B1 = basis_map(B1_type);
Nb1 = B1.Nb();
Nlb1 = B1.Nlb();
B2 = basis_map(B2_type);
Nb2 = B2.Nb();
Nlb2 = B2.Nlb();
u_map = containers.Map(...
    {'u1', 'u2'},...
    {struct('shift', 0),...
    struct('shift', mesh.Nb)});
U = u_map(u_type);
shift = U.shift;
rows = zeros(2*Nlb1*Nlb2,1);
cols = rows;
values = rows;
flag = 1;
for n = 1:mesh.Ne
    E = mesh.node(mesh.elem(n,:),:);
    V = E(1:3,:);
    [gauss_points_loc, gauss_weights_loc] = gauss_lengdre_loc(V, "2d");
    phiu = basis_function_loc(V, gauss_points_loc(:,1), gauss_points_loc(:,2), "P1", du_type);
    phi1 = basis_function_loc(V, gauss_points_loc(:,1), gauss_points_loc(:,2), B1_type, D1_type);
    phi2 = basis_function_loc(V, gauss_points_loc(:,1), gauss_points_loc(:,2), B2_type, D2_type);
    item = phiu * U_old(mesh.elem(n,1:3)+shift);
    for i = 1:Nlb2
        for j = 1:Nlb1
            f = item .* phi1(:,j) .* phi2(:,i);
            r = gauss_weights_loc * f;
            rows(flag) = mesh.elem(n,i);
            cols(flag) = mesh.elem(n,j);
            values(flag) = r;
            flag = flag + 1;
        end
    end
end
matrix_A = sparse(rows, cols, values, Nb2, Nb1);
end

%% assemble_nonlinear_vector
function vector_b = assemble_nonlinear_vector(mesh, u1_type, du1_type, u2_type, du2_type, B_type, D_type, U_old)
u_map = containers.Map(...
    {'u1', 'u2'},...
    {struct('shift', 0),...
    struct('shift', mesh.Nb)});
U1 = u_map(u1_type);
shift1 = U1.shift;
U2 = u_map(u2_type);
shift2 = U2.shift;
vector_b = zeros(mesh.Nb,1);
for n = 1:mesh.Ne
    E = mesh.node(mesh.elem(n,:),:);
    V = E(1:3,:);
    [gauss_points_loc, gauss_weights_loc] = gauss_lengdre_loc(V, "2d");
    phi1 = basis_function_loc(V, gauss_points_loc(:,1), gauss_points_loc(:,2), "P1", du1_type);
    phi2 = basis_function_loc(V, gauss_points_loc(:,1), gauss_points_loc(:,2), "P1", du2_type);
    phi3 = basis_function_loc(V, gauss_points_loc(:,1), gauss_points_loc(:,2), B_type, D_type);
    item1 = phi1 * U_old(mesh.elem(n,1:3)+shift1);
    item2 = phi2 * U_old(mesh.elem(n,1:3)+shift2);
    for j = 1:mesh.Nlb
        f =  item1 .* item2 .* phi3(:,j);
        r = gauss_weights_loc * f;
        vector_b(mesh.elem(n,j),1) = vector_b(mesh.elem(n,j),1) + r;
    end
end
end